import time
from datetime import timedelta, timezone, datetime
import argparse

from yms_zsl.train import (
    train_feature_extractor,
    train_decae,
    semantics,
    train_fcnn,
    predict
)


def main(args):
    # 时间初始化
    run = wandb.init(project='ZSL', name='HZL', id='001')
    start_time = datetime.now(timezone(timedelta(hours=8))).strftime("%Y-%m-%d %H:%M:%S")
    t1 = time.time()
    print(f"程序于 {start_time} 开始运行!!!\n")

    # 参数解析
    data = args.data
    data_root = args.data_root
    output_root = args.output_root

    # --------------------------- 训练特征提取器 ---------------------------
    feature_opts = train_feature_extractor.parse_args([
        "--data_dir", f"{data_root}/{data}",
        "--save_dir", f"{output_root}/train_{data}",
        "--epochs", "50",
        "--lr", "0.001",
        "--batch_size", "128",
        "--num_classes", "10"
    ])
    run = train_feature_extractor.main(feature_opts, run)

    t2 = time.time()
    print(f"训练特征提取器用时: {str(timedelta(seconds=int(t2 - t1)))}")

    # --------------------------- 训练编码器 ---------------------------
    decae_opts = train_decae.parse_args([
        "--data_dir", f"{data_root}/{data}",
        "--save_dir", f"{output_root}/train_{data}",
        "--epochs", "100",
        "--lr", "0.001",
        "--batch_size", "128"
    ])
    run = train_decae.main(decae_opts, run)

    t3 = time.time()
    print(f"训练编码器用时: {str(timedelta(seconds=int(t3 - t2)))}s")

    # --------------------------- 提取语义信息 ---------------------------
    semantics_opts = semantics.parse_args([
        "--data_dir", f"{data_root}/{data}",
        "--encoder_path", f"{output_root}/train_{data}/models/decae.pt",
        "--save_path", f"{output_root}/train_{data}/HSA.mat",
        "--feature_extractor_path", f"{output_root}/train_{data}/models/feature_extractor.pt",
        "--features", f"{output_root}/train_{data}/{data}-{data}.mat",
        "--batch_size", "100",
        "--sa_path", ""
    ])
    semantics.semantics(semantics_opts)

    t4 = time.time()
    print(f"提取用时: {str(timedelta(seconds=int(t4 - t3)))}s")

    # --------------------------- 训练 FCNN ---------------------------
    # fcnn_opts = train_fcnn.parse_args([
    #     "--data_dir", f"{data_root}/{data}",
    #     "--output_dir", f"{output_root}/train_{data}",
    #     "--feature_path", f"{output_root}/train_{data}/models/feature_extractor.pt",
    #     "--hsa_path", f"{output_root}/train_{data}/HSA.mat",
    #     "--lr", "1e-3",
    #     "--batch_size", "64",
    #     "--epochs", "50",
    #     "--fcnn_channels", "518"  # 修复了多余的数字 56（原为无效参数）
    # ])
    # train_fcnn.main(fcnn_opts)
    #
    # t5 = time.time()
    # print(f"训练 FCNN 用时: {str(timedelta(seconds=int(t5 - t4)))}s")
    #
    # # --------------------------- 预测验证 ---------------------------
    # predict_opts = predict.parse_args([
    #     "--fcnn_path", f"{output_root}/train_{data}/models/best_fcnn.pt",
    #     "--data_path", f"{data_root}",
    #     "--hsa_path", f"{output_root}/train_{data}/HSA.mat",
    #     "--feature_path", f"{output_root}/train_{data}/models/feature_extractor.pt",
    #     "--cnn_path", f"{output_root}/train_{data}/models/best_cnn.pt",
    #     "--output_dir", f"{output_root}/train_{data}/predict",
    #     "--batch_size", "32"
    # ])
    # predict.main(predict_opts)
    #
    # t6 = time.time()
    print(f"验证用时: {str(timedelta(seconds=int(t4 - t3)))}s")
    print(f"整体用时: {str(timedelta(seconds=int(t4 - t1)))}s")
    run.finish()


if __name__ == "__main__":
    import wandb

    wandb.login(key='7f07d023ba84876d3a2359818094ab3f427f9794')
    parser = argparse.ArgumentParser(description="训练和预测脚本")
    parser.add_argument("--data", type=str, help="数据标识", default="D0")
    parser.add_argument('--output_root', type=str, default='/data/coding/output')
    parser.add_argument('--data_root', type=str, default="/data/coding/CRWU")
    opts = parser.parse_args()
    main(opts)
